#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

import hashlib
from collections import defaultdict

from Bio import SeqIO
from Bio.SeqRecord import SeqRecord
import numpy as np
from tqdm import tqdm
import sys
import os


def get_md5_value(ss: str):
    my_md5 = hashlib.md5()  # 获取一个MD5的加密算法对象
    my_md5.update(ss.encode("utf-8"))  # 得到MD5消息摘要
    my_md5_Digest = my_md5.hexdigest()  # 以16进制返回消息摘要，32位
    return my_md5_Digest


def check_genbank_unique(filepath: str):
    dir_name = os.path.dirname(filepath)
    file_name = os.path.basename(filepath)
    file_name = file_name[:file_name.rindex(".")]
    ht = defaultdict(list)
    for seq_record in tqdm(SeqIO.parse(filepath, "fasta"), desc="seqs"):
        fasta = seq_record.seq.upper()
        seq_id = seq_record.id
        md5 = get_md5_value(fasta)
        ht[md5].append(seq_id)
    gen_report(ht, os.path.join(dir_name, file_name + ".report"))
    total_fasta = []
    for seq_record in tqdm(SeqIO.parse(filepath, "fasta"), desc="seqs"):
        fasta = seq_record.seq.upper()
        seq_id = seq_record.id
        md5 = get_md5_value(fasta)
        if md5 in ht:
            seq_id, seq_des = manage_title(ht[md5])
            new_record = SeqRecord(fasta, id=seq_id, description=seq_des)
            total_fasta.append(new_record)
            ht.pop(md5)
    SeqIO.write(total_fasta, os.path.join(dir_name, file_name + ".unique.fasta"), "fasta")


def manage_title(seq_ids: list):
    title = ""
    if len(seq_ids) == 1:
        return seq_ids[0], ""
    else:
        for i in range(0, len(seq_ids)):
            if seq_ids[i].startswith("NC_") or seq_ids[i].startswith("AC_"):
                title = seq_ids[i]
                seq_ids.pop(i)
                break
        if title == "":
            title = seq_ids[0]
            seq_ids.pop(0)
        description = "|".join(seq_ids)
        return title, description


def gen_report(ht: defaultdict, outpath: str):
    unique_seq = len(ht)
    total_seq = 0
    seqs = []
    for key in ht.keys():
        length = len(ht[key])
        total_seq += length
        seqs.append(length)
    seqs.sort()
    min_seq = seqs[0]
    max_seq = seqs[-1]
    mean_seq = np.mean(seqs)
    median_seq = np.median(seqs)
    outfile = open(outpath, "w")
    outfile.write("Total sequences = \t" + str(total_seq) + "\n")
    outfile.write("Unique sequences = \t" + str(unique_seq) + "\n")
    outfile.write("Max sequences = \t" + str(max_seq) + "\n")
    outfile.write("Min sequences = \t" + str(min_seq) + "\n")
    outfile.write("Mean sequences = \t" + str(mean_seq) + "\n")
    outfile.write("Median sequences = \t" + str(median_seq) + "\n")
    outfile.flush()
    outfile.close()


if __name__ == '__main__':
    check_genbank_unique(sys.argv[1])
